﻿#pragma once

#include <iomanip>
#include <iostream>
#include <map>
#include <ostream>
#include <string>
#include <vector>
#include <yaml-cpp/yaml.h>

#include "aimodel/model_error_code.h"
#include "aimodel/model_struct.h"
#include "common/transforms.h"
#include "base_util/utils.h"

class LogInfo;

namespace ai {

class ModelConfig {
public:
  ModelConfig(LogInfo* li):log_ifo(li) {
    algorithm_map["cls"] = model_cls;
    algorithm_map["cls_video"] = model_cls_video;
    algorithm_map["text_det"] = model_text_det;
    algorithm_map["text_rec"] = model_text_rec;
    algorithm_map["seg"] = model_seg;
    algorithm_map["det"] = model_det;
    algorithm_map["plate_rec"] = model_plate_rec;
    algorithm_map["face_det"] = model_face_det;
    algorithm_map["face_rec"] = model_face_rec;
  }
  ~ModelConfig(){}

  void reset() {
    use_gpu = 0;
    gpu_ids.clear();
    gpu_mem = 4000;
    predictor_type = "";
    label_list = {};
  }

  int load_config(const std::string& config_path) {
    reset();
    try {
      config = YAML::LoadFile(config_path);
    } catch(...) {
      std::cerr << "YAML::LoadFile(config_path) error:" << config_path << std::endl;
      return model_load_config_error;
    }
    
    transforms = new Transforms(log_ifo);
    int ret = transforms->init(config["Transforms"]);
    if (ret != 0) {return model_load_config_transforms_error;}
    
    // printf("transforms->init sucess\n");
    if (config["algorithm"].IsDefined()) {
      algorithm_str = config["algorithm"].as<std::string>();
      algorithm = algorithm_map[this->algorithm_str];
    } else {
      std::cerr << "Please set algorithm: cls/det/seg";
      return model_load_config_error;
    }

    if (config["use_gpu"].IsDefined()) { use_gpu = config["use_gpu"].as<int>();} 
    else {use_gpu = 0;}

    if (config["gpu_ids"].IsDefined()) { gpu_ids = config["gpu_ids"].as<std::vector<int>>();} 
    else {gpu_ids = {0};}

    if (config["gpu_mem"].IsDefined()) { gpu_mem = config["gpu_mem"].as<int>();} 
    else {gpu_mem = 500;}

    if (config["cpu_threads"].IsDefined()) { cpu_threads = config["cpu_threads"].as<int>();} 
    else {cpu_threads = 2;}

    if (config["use_mkldnn"].IsDefined()) { use_mkldnn = config["use_mkldnn"].as<bool>();} 
    else {use_mkldnn = false;}

    if (config["image_type"].IsDefined()) { img_type = config["image_type"].as<std::string>();} 
    else { img_type = "RGB"; }      // BGR

    if (config["channels"].IsDefined()) { channels = config["channels"].as<int>(); } 
    else { channels = 3; }

    if (config["data_format"].IsDefined()) { data_format = config["data_format"].as<std::string>(); } 
    else { data_format = "CHW"; }   // HWC

    if (config["arch"].IsDefined()) { arch = config["arch"].as<std::string>();} 
    else { arch = "unknow"; }    // resnet, yolo, ssd, maskrcnn, unet...
    if (config["min_subgraph_size"].IsDefined()) { min_subgraph_size = config["min_subgraph_size"].as<int>(); } 
    else { min_subgraph_size = 3; }
    if (config["draw_threshold"].IsDefined()) { draw_threshold = config["draw_threshold"].as<float>(); } 
    else { draw_threshold = 0.5; }
    if (config["nms_threshold"].IsDefined()) { nms_threshold = config["nms_threshold"].as<float>(); } 
    else { nms_threshold = 0.5; }
    if (config["input_shape"].IsDefined()) { input_shape = config["input_shape"].as<std::vector<int>>(); } 
    else { input_shape = {640, 640}; }
    if (config["strides"].IsDefined()) { strides = config["strides"].as<std::vector<int>>(); } 
    else { strides = {8, 16, 32}; }
    if (config["anchors_table"].IsDefined()) { anchors_table = config["anchors_table"].as<std::vector<std::vector<std::vector<int>>>>(); } 
    else { anchors_table = {}; }
    spdlog::get("logger")->info("anchors_table: {}", anchors_table.size());
    if (config["map_size"].IsDefined()) { map_size = config["map_size"].as<std::vector<std::vector<int>>>(); } 

    if (config["predictor_type"].IsDefined()) { predictor_type = config["predictor_type"].as<std::string>();} 
    else { predictor_type = "analysis"; }    // analysis, fp16, fp32, int8

    if (config["batch_size"].IsDefined()) { batch_size = config["batch_size"].as<int>(); } 
    else { batch_size = 1; }

    if (config["input_names"].IsDefined()) {
      input_names = config["input_names"].as<std::vector<std::string>>();
    } 
    if (config["output_names"].IsDefined()) {
      output_names = config["output_names"].as<std::vector<std::string>>();
    } 
    if (config["label_list"].IsDefined()) {
      label_list = config["label_list"].as<std::vector<std::string>>();
    } 
    else if (algorithm == model_text_rec && config["label_path"].IsDefined()) {
      label_list = ReadDict(bf::path(config_path).parent_path().append(config["label_path"].as<std::string>()).string());
    } 
    else if (algorithm == model_text_det || algorithm == model_plate_rec) {
      // pass
    } else {
      spdlog::get("logger")->info("Please set label_list.");
      return model_load_config_error;
    }

    // model_text_det
    if (config["max_side_len"].IsDefined()) { max_side_len = config["max_side_len"].as<int>(); } 
    if (config["det_db_thresh"].IsDefined()) { det_db_thresh = config["det_db_thresh"].as<float>(); } 
    if (config["det_db_box_thresh"].IsDefined()) { det_db_box_thresh = config["det_db_box_thresh"].as<float>(); } 
    if (config["det_db_unclip_ratio"].IsDefined()) { det_db_unclip_ratio = config["det_db_unclip_ratio"].as<float>(); } 
    if (config["det_expand_by_min_score"].IsDefined()) { det_expand_by_min_score = config["det_expand_by_min_score"].as<float>(); } 
    if (config["use_space_char"].IsDefined()) { use_space_char = config["use_space_char"].as<bool>();} 

    printf("load_config done\n");
    return 0;
  }

  void print_config_info() const {
    spdlog::get("logger")->info("=======Paddle {} inference config======", algorithm_str);
    spdlog::get("logger")->info("use_gpu: {}", use_gpu);
    spdlog::get("logger")->info("image_type: {}", img_type);
    spdlog::get("logger")->info("channels: {}", channels);
    spdlog::get("logger")->info("predictor_type: {}", predictor_type);
    spdlog::get("logger")->info("batch_size: {}", batch_size);

    spdlog::get("logger")->info("label_list: [");
    for (size_t i = 0; i < label_list.size(); ++i) {
        if (i != label_list.size() - 1) { spdlog::get("logger")->info("{}, ", label_list[i]); } 
        else { spdlog::get("logger")->info(label_list[i]); }
    }
    spdlog::get("logger")->info("]");
    spdlog::get("logger")->info("=======End of Paddle {} inference config======", algorithm_str);
    }

public:
  bool use_gpu;
  int gpu_id;
  std::vector<int> gpu_ids;
  int gpu_mem;
  bool use_mkldnn;
  int cpu_threads;

  std::vector<int> input_shape;
  std::vector<int> strides;
  std::vector<std::vector<std::vector<int>>> anchors_table;
  std::vector<std::vector<int>> map_size;

  std::string algorithm_str;        // 检测算法名称
  ModelType algorithm;
  std::map<std::string, ModelType> algorithm_map;

  std::string arch;
  int min_subgraph_size;
  float draw_threshold;
  float nms_threshold;

  std::string predictor_type;
  std::string img_type;
  int channels;
  std::string data_format;
  int batch_size;
  std::vector<std::string> input_names;
  std::vector<std::string> output_names;
  std::vector<std::string> label_list;

  // text det
  int use_space_char = false;
  int max_side_len = 960;
  double det_db_thresh = 0.3;
  double det_db_box_thresh = 0.5;
  double det_db_unclip_ratio = 2.0;
  double det_expand_by_min_score = 0.1;

  Transforms* transforms;
  YAML::Node config;
  LogInfo* log_ifo;
};

} // namespace ai
